#!/usr/bin/env python3
"""
Stores jobs in a MongoDB database.
"""

import logging

from apscheduler.jobstores.base import JobStore
from apscheduler.job import Job

try:
  import pickle as pickle
except ImportError:  # pragma: nocover
  import pickle

try:
  from bson.binary import Binary
  from pymongo.connection import Connection
except ImportError:  # pragma: nocover
  raise ImportError("MongoDBJobStore requires PyMongo installed")

logger = logging.getLogger(__name__)


class MongoDBJobStore(JobStore):
  def __init__(
    self,
    database="apscheduler",
    collection="jobs",
    connection=None,
    pickle_protocol=pickle.HIGHEST_PROTOCOL,
    **connect_args,
  ):
    self.jobs = []
    self.pickle_protocol = pickle_protocol

    if not database:
      raise ValueError('The "database" parameter must not be empty')
    if not collection:
      raise ValueError('The "collection" parameter must not be empty')

    if connection:
      self.connection = connection
    else:
      self.connection = Connection(**connect_args)

    self.collection = self.connection[database][collection]

  def add_job(self, job):
    job_dict = job.__getstate__()
    job_dict["trigger"] = Binary(pickle.dumps(job.trigger, self.pickle_protocol))
    job_dict["args"] = Binary(pickle.dumps(job.args, self.pickle_protocol))
    job_dict["kwargs"] = Binary(pickle.dumps(job.kwargs, self.pickle_protocol))
    job.id = self.collection.insert(job_dict)
    self.jobs.append(job)

  def remove_job(self, job):
    self.collection.remove(job.id)
    self.jobs.remove(job)

  def load_jobs(self):
    jobs = []
    for job_dict in self.collection.find():
      try:
        job = Job.__new__(Job)
        job_dict["id"] = job_dict.pop("_id")
        job_dict["trigger"] = pickle.loads(job_dict["trigger"])
        job_dict["args"] = pickle.loads(job_dict["args"])
        job_dict["kwargs"] = pickle.loads(job_dict["kwargs"])
        job.__setstate__(job_dict)
        jobs.append(job)
      except Exception:
        job_name = job_dict.get("name", "(unknown)")
        logger.exception('Unable to restore job "%s"', job_name)
    self.jobs = jobs

  def update_job(self, job):
    spec = {"_id": job.id}
    document = {"$set": {"next_run_time": job.next_run_time}, "$inc": {"runs": 1}}
    self.collection.update(spec, document)

  def close(self):
    self.connection.disconnect()

  def __repr__(self):
    connection = self.collection.database.connection
    return f"<{self.__class__.__name__} (connection={connection})>"
